import math
import torch
import utils
from torch.nn.modules.module import Module
from torch.nn.parameter import Parameter
from torch import nn


class GraphConvolution(nn.Module):
    def __init__(self, input_dim):
        super(GraphConvolution, self).__init__()
        self.node_theta = Parameter(torch.FloatTensor(1, 2 * input_dim))
        self.path_theta = Parameter(torch.FloatTensor(1, 2 * input_dim))
        self.net_theta = Parameter(torch.FloatTensor(1, 2 * input_dim))
        self.reset_parameters()

    def save_weight(self):
        pass

    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.node_theta.size(1))
        self.node_theta.data.uniform_(-stdv, stdv)
        self.path_theta.data.uniform_(-stdv, stdv)
        self.net_theta.data.uniform_(-stdv, stdv)

    def forward(self, W, H, n, path_threshold):
        node_H = utils.find_n_node(W, H, self.node_theta, n)
        path, path_H = utils.find_n_path(n, W, H, self.path_theta, path_threshold)
        net, net_H, net_adj = utils.find_n_network(n, W, H, self.net_theta)
        return node_H, path_H, net_H

    def __repr__(self):
        return self.__class__.__name__ + ' (' \
               + str(self.input_dim) + ' -> ' \
               + str(self.output_dim) + ')'
